# 提取ROS_BAG数据
import rosbag
import cv2 as cv
from cv_bridge import CvBridge
import sensor_msgs.point_cloud2 as pc2
import numpy as np

# 解析ROS-BAG路径
bag_dir = '2023-06-21-12-45-04.bag'
bag = rosbag.Bag(bag_dir)
data_types, topics = bag.get_type_and_topic_info()
for d_type in data_types.keys():
    print(d_type, end=': ')
    print(data_types[d_type])
for topic in topics.keys():
    print(topic, end=': ')
    print(topics[topic])
img_topic = '/src_rgb'
pcd_topic = '/rslidar_points'
# 解析数据帧数
max_num = 20

num_data = 0
bridge = CvBridge()
for topic, msg, t_stamp in bag.read_messages(topics=[img_topic]):
    data = bridge.imgmsg_to_cv2(msg)
    cv.imwrite('tmp_img/tmp_' + str(num_data) + '.png', data)
    cv.imshow('new', cv.resize(data, [960, 540]))
    cv.waitKey(100)
    num_data += 1
    if num_data >= 3 * max_num:
        break

num_data = 0
for topic, msg, t_stamp in bag.read_messages(topics=[pcd_topic]):
    data = pc2.read_points(msg, skip_nans=True, field_names=('x', 'y', 'z'))
    data = np.array(list(data))
    np.savetxt('tmp_pcd/tmp_' + str(num_data) + '.txt', data, encoding='utf-8')
    np.save('tmp_pcd/tmp_' + str(num_data), data)
    num_data += 1
    if num_data >= max_num:
        break
